'''
经常会用到的一些函数
'''
import torch

def cos_sim(mat1:torch.Tensor, mat2:torch.Tensor) -> torch.Tensor:
    # 计算点乘矩阵 - 分子
    numerator = torch.matmul(mat1, mat2.T)
    # 计算点乘矩阵 - 分母
    denominator = torch.matmul(
        torch.sqrt(torch.sum(mat1**2, dim=-1)).unsqueeze(-1),
        torch.sqrt(torch.sum(mat2**2, dim=-1)).unsqueeze(-2)
    )
    return numerator / denominator

def seq_to_adjacency(seq:str, sub_index:dict, sub_len:int = 3, gap:int = 0) -> torch.Tensor:
    '''
    seq:需要转换为图的碱基序列
    sub_index:碱基序列子串的索引,便于在矩阵中定位碱基子串
    sub_len:子串长度
    gap:划分子串的间隔
    return:一个二维整数矩阵(有向图邻接矩阵),可能包含途图中不存在的点。即碱基序列可能不包含所有可能的子串
    '''
    # 需要返回的图
    re = torch.zeros((len(sub_index), len(sub_index)), dtype=torch.float32)

    # 注：用range(len(seq))时，i=0与i=sub_len + edge_gap截取到的位置是一样的。
    for begin_index in  range(sub_len + gap):#获取第一个子串的开始位置。
        sub_str = seq[begin_index:begin_index + sub_len]#获取图的一个点
        if len(sub_str) != sub_len:
            continue
        # re[sub_index[sub_str]][sub_index[sub_str]] += 1# 标记点。点出现的数量是否有意义？
        # 标记边
        sub_str_pre = sub_str#记录已经遍历的点，下次直接从下一个点遍历
        for j in range(begin_index + sub_len + gap, len(seq) - 1, sub_len + gap):
            sub_str = seq[j:j + sub_len]#获取图的下一个点
            if len(sub_str) != sub_len:
                continue
            # 标记边(有向)，边的数量有无意义？
            # re[sub_index[sub_str]][sub_index[sub_str_pre]] += 1
            re[sub_index[sub_str_pre]][sub_index[sub_str]] = 1
            sub_str_pre = sub_str# 更新遍历过的点
    return re
